import torch
import torch.nn as nn
import time

"""
    VGG implemented by PyTorch
    Website: https://arxiv.org/pdf/1409.1556v6.pdf
"""

cfg = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
}


class VGG(nn.Module):

    def __init__(self, features, num_class=100):
        super().__init__()
        self.features = features

        # three Fully-Connected Layers
        self.classifier = nn.Sequential(
            nn.Linear(512, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, num_class)
        )

    def forward(self, x):
        output = self.features(x)
        output = output.view(output.size(0), -1)
        output = self.classifier(output)

        return output

    @staticmethod
    def make_layers(cfg, batch_norm=False):
        layers = []

        input_channel = 3
        for config in cfg:
            if config == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(input_channel, config, kernel_size=3, padding=1)]

                if batch_norm:
                    layers += [nn.BatchNorm2d(config)]

                layers += [nn.ReLU(inplace=True)]

                input_channel = config

        return nn.Sequential(*layers)


def vgg11_bn():
    return VGG(VGG.make_layers(cfg['A'], batch_norm=True))


def vgg13_bn():
    return VGG(VGG.make_layers(cfg['B'], batch_norm=True))


def vgg16_bn():
    return VGG(VGG.make_layers(cfg['D'], batch_norm=True))


def vgg19_bn():
    return VGG(VGG.make_layers(cfg['E'], batch_norm=True))


if __name__ == '__main__':
    a = time.time()
    m = nn.Linear(2000, 30)
    m.device = 'cuda'
    input = torch.randn(128, 2000)
    output = m(input)
    pass
    print(output.shape)
    print(time.time() - a)
